# 绘制训练结果

import json
import configparser

import matplotlib.pyplot as plt
import os
import src.tool as tool
import numpy as np

# 配置文件地址
filename = './config.ini'
# 读取程序配置
config = configparser.ConfigParser()
config.read(filename, encoding='utf-8')

LR = tool.get_config_by_key(config['ai'], 'LR')
BATCH_SIZE = tool.get_config_by_key(config['ai'], 'BATCH_SIZE')
filename = './saved_model/train_data.txt'


def paint(all_data: list):
    epoch_winner = []
    loss = []
    accuracy = []
    epoch = 0
    loss_mean = []
    for i in range(len(all_data)):
        single = all_data[i]
        try:
            parsed = json.loads(single)
        except Exception as e:
            print("第{}行不是需要的数据，错误原因->{}".format(i + 1, e))
            continue
        data_type = parsed['type']
        data = parsed['data']
        if data_type == 'loss':
            loss.extend(data)
            loss_mean.append(np.mean(data))
            epoch = epoch + 1
        elif data_type == 'epoch_winner':
            epoch_winner.extend(data)
        elif data_type == 'accuracy':
            for j in range(len(data)):
                single = data[j]
                percent = single['percent']
                accuracy.append(percent)

    plt.plot(loss)
    plt.ylabel('all_loss')
    plt.xlabel('LR={},batch_size = {},epoch = {}'.format(LR, BATCH_SIZE, epoch))
    plt.show()
    plt.plot(loss_mean)
    plt.ylabel('loss_mean')
    plt.xlabel('LR={},batch_size = {},epoch = {}'.format(LR, BATCH_SIZE, epoch))
    plt.show()
    plt.plot(epoch_winner)
    plt.ylabel('all_winner')
    plt.xlabel('LR={},batch_size = {},epoch = {}'.format(LR, BATCH_SIZE, epoch))
    plt.show()
    plt.plot(accuracy)
    plt.ylabel('all_accuracy')
    plt.xlabel('LR={},batch_size = {},epoch = {}'.format(LR, BATCH_SIZE, epoch))
    plt.show()

    print('all_loss->',loss_mean)
    decrease = []
    # 损失下降幅度
    for i in range(len(loss_mean) - 1):
        a = loss_mean[i]
        b = loss_mean[i+1]
        decrease.append(a-b)

    plt.plot(decrease)
    plt.ylabel('decrease')
    plt.xlabel('LR={},batch_size = {},epoch = {}'.format(LR, BATCH_SIZE, epoch))
    plt.show()


if __name__ == '__main__':
    fs = open(filename, encoding='utf-8', mode='r')
    all_data = fs.read()
    lines = all_data.split('\n')
    paint(lines)
